package com.shujia;

import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.util.ArrayList;
import java.util.List;


/**
 * 需求：
 *      根据字符串{"movie": [{"movie_name": "肖申克的救赎", "MovieType": "犯罪" }, {"movie_name": "肖申克的救赎", "MovieType": "剧情" }]}
 *      获取其中的movie_name对应的value作为一列，MovieType作为另外一列，其中JSON数组中每一个JSON都是一行数据
 *
 *   select get_json_object('{"movie": [{"movie_name": "肖申克的救赎", "MovieType": "犯罪" }, {"movie_name": "肖申克的救赎", "MovieType": "剧情" }]}','$.movie');
 *
 *  分析：
 *      根据get_json_object获取其中的movieKey对应的JSONArray,并将该结果传入UDTF函数进行解析，返回对行数据
 *
 *  解析过程：
 *      根据 [{"movie_name":"肖申克的救赎","MovieType":"犯罪"},{"movie_name":"肖申克的救赎","MovieType":"剧情"}]
 *      需要将字符串转换成JSONArray对象，然后循环取movie_name和 MovieType，之后通过forward函数将一行数据写出
 *
 *  建表语句：
 *      CREATE TABLE movie_json(movie string);
 *
 *  加载数据：
 *      load data local inpath '/usr/local/soft/data/UDTF.txt' into table movie_json;
 *
 *  测试表：
 *      select get_json_object(movie,'$.movie') from movie_json;
 *
 *  add jar /usr/local/soft/jars/hiveCode-1.0.jar
 *
 *  create temporary function getjson as 'com.shujia.MyUDTF2';
 *
 *  select getjson(get_json_object(movie,'$.movie')) from movie_json;
 *
 */
public class MyUDTF2 extends GenericUDTF {
    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentLengthException, UDFArgumentTypeException {
        // 通过StructObjectInspector获取所有参数，其中StructField表示对应一个参数
        List<? extends StructField> allStructFieldRefs = argOIs.getAllStructFieldRefs();

        // 判断参数个数
        if(allStructFieldRefs.size() != 1){
            throw new UDFArgumentLengthException("该函数参数为1个JSONArray的字符串");
        }

        // 判断参数类型
        StructField structField = allStructFieldRefs.get(0);
        ObjectInspector objectInspector = structField.getFieldObjectInspector();
        if(!objectInspector.getCategory().equals(ObjectInspector.Category.PRIMITIVE)){
            throw new UDFArgumentTypeException(1,"参数的类型不正确，请使用HIVE中的普通数据类型String");
        }

        // 返回字段名称
        ArrayList<String> fieldNames = new ArrayList<String>();
        fieldNames.add("movieName");
        fieldNames.add("movieType");


        // 返回字段的类型
        ArrayList<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>();
        // 多字段时，添加字段类型顺序需要和字段名称顺序保持一致
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);


        return  ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames,
                fieldOIs);
    }

    /**
     * 用于处理JSONArray字符串
     * @param args [{"movie_name":"肖申克的救赎","MovieType":"犯罪"},{"movie_name":"肖申克的救赎","MovieType":"剧情"}]
     * @throws HiveException
     */
    @Override
    public void process(Object[] args) throws HiveException {
        String jsonStr = args[0].toString();
        try {
            // 将字符串加载成JSONArray对象
            JSONArray jsonArray = new JSONArray(jsonStr);

            // JSONObject
//            JSONObject jsonObject1 = new JSONObject(jsonStr);

            String[] line_data = new String[2];

            // 通过获取长度，进行遍历
            for (int i = 0; i < jsonArray.length(); i++) {
                // 通过getJSONObject获取jsonArray中每一个JSON对象：{"movie_name":"肖申克的救赎","MovieType":"犯罪"}
                JSONObject jsonObject = jsonArray.getJSONObject(i);
                String movie_name = jsonObject.getString("movie_name");
                String movie_type = jsonObject.getString("MovieType");
                // 需要和fieldNames中定义字段名称顺序一致
                line_data[0] = movie_name;
                line_data[1] = movie_type;
                // 获取数据后，将结果返回
                forward(line_data);
            }
            
        } catch (JSONException e) {
            e.printStackTrace();
        }

    }

    @Override
    public void close() throws HiveException {
        // TODO: 2022/9/15 nothing todo 
    }
}

/**
 * 作业:
 *      创建一个UDTF函数，实现功能为 传入一个JSON字符串 {"movie": [{"movie_name": "肖申克的救赎", "MovieType": "犯罪" }, {"movie_name": "肖申克的救赎", "MovieType": "剧情" }]}
 *          解析其中的movie_name和MovieType转换为两列多行的结果
 *
 */